import csv
import pymysql
import argparse


def create_table_and_insert_data(csv_file, db_config, batch_size=1000):
    # 连接到MySQL数据库
    conn = pymysql.connect(**db_config, charset='utf8mb4')
    cursor = conn.cursor()

    # 创建数据库和表（如果不存在的话）
    cursor.execute("CREATE DATABASE IF NOT EXISTS blg CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci")
    cursor.execute("USE blg")

    # 创建表
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS find_file (
            id INT AUTO_INCREMENT PRIMARY KEY,
            file_name VARCHAR(255),
            absolute_path TEXT,
            file_extension VARCHAR(50),
            file_size BIGINT,
            md5 VARCHAR(32)
        ) ENGINE=InnoDB CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci
    ''')

    # 创建索引
    cursor.execute("CREATE INDEX idx_file_ext_size ON find_file (file_extension, file_size)")

    # 读取CSV文件并插入数据
    with open(csv_file, mode='r', encoding='utf-8') as csvfile:
        reader = csv.DictReader(csvfile)
        batch = []
        count = 0
        for row in reader:
            batch.append((
                row['File Name'],
                row['Absolute Path'],
                row['File Extension'],
                row['File Size (bytes)'],
                row['MD5 Hash']
            ))
            count += 1

            # 每处理完 batch_size 条数据，就批量插入到数据库
            if count % batch_size == 0:
                print(f"Processing batch {count // batch_size} with {batch_size} records.")
                cursor.executemany('''
                    INSERT INTO find_file (file_name, absolute_path, file_extension, file_size, md5)
                    VALUES (%s, %s, %s, %s, %s)
                ''', batch)
                conn.commit()  # 提交当前批次
                batch = []  # 清空批次

        # 处理最后剩余的记录（如果有的话）
        if batch:
            print(f"Processing batch {count // batch_size + 1} with {len(batch)} records.")
            cursor.executemany('''
                INSERT INTO find_file (file_name, absolute_path, file_extension, file_size, md5)
                VALUES (%s, %s, %s, %s, %s)
            ''', batch)
            conn.commit()

    # 提交更改并关闭连接
    cursor.close()
    conn.close()


if __name__ == "__main__":
    # 设置命令行参数解析
    parser = argparse.ArgumentParser(description='根据CSV文件创建MySQL数据库表并插入数据')
    parser.add_argument('--csv_file', type=str, help='输入的CSV文件路径')
    parser.add_argument('--host', type=str, help='MySQL主机地址')
    parser.add_argument('--user', type=str, help='MySQL用户名')
    parser.add_argument('--password', type=str, help='MySQL密码')

    # 解析命令行参数
    args = parser.parse_args()

    # MySQL数据库配置
    db_config = {
        'host': args.host,
        'user': args.user,
        'password': args.password,
        'database': 'blg'
    }

    # 调用函数
    create_table_and_insert_data(args.csv_file, db_config)
